import os
from PIL import Image


image_folder = r'./data/neu_seg_competition_nation_self/images/test'
mask_folder = r'./data/neu_seg_competition_nation_self/annotations/test'
image_folder_1 = r'./data/neu_seg_competition_nation_self/images/training'
mask_folder_1 = r'./data/neu_seg_competition_nation_self/annotations/training'

# 检测并转换图像通道
def convert_to_single_channel(folder_path):
    for filename in os.listdir(folder_path):
        if filename.endswith(('.png', '.jpg')):
            file_path = os.path.join(folder_path, filename)
            # print(file_path)
            img = Image.open(file_path)
            if img.mode != 'L':  # 检查是否为单通道图像
                img = img.convert('L')  # 转换为单通道图像
                img.save(file_path)  # 保存图像
                print(f"Converted {filename} to single channel.")

# 处理测试集和训练集图像
convert_to_single_channel(image_folder)
convert_to_single_channel(mask_folder)
convert_to_single_channel(image_folder_1)
convert_to_single_channel(mask_folder_1)